import os.path

import numpy as np
import torch
from torch import nn, optim
import matplotlib.pyplot as plt
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
import torch.nn.functional as F


# 先不用torch的线性回归api
def test1():
    learning_rate = 0.01  # 学习率

    # 1、准备数据
    # 公式：y=3x+0.8
    x = torch.rand(500, 1)  # 500行1列
    y_true = x * 3 + 0.8  # 这是结果，w=3，b=0.8，通过梯度降低求出最接近这个正确答案的值。真实场景中是不知道这个公式的，只有y的值。
    # 2、通过模型计算y_predict
    w = torch.rand(1, 1, requires_grad=True)
    b = torch.tensor(0, requires_grad=True, dtype=torch.float32)

    # 4、通过循环，反向传播，更新参数
    for i in range(500):
        y_predict = torch.matmul(x, w) + b  # 这是预测值，y=w*x+b
        # 3、计算loss
        loss = (y_true - y_predict).pow(2).mean()  # 损失=（真实值-预测值）的平方即方差的平均值

        if w.grad is not None:
            w.data.zero_()  # 每次都清零，这样防止因为旧的记录导致重复计算，每次循环都是独立计算梯度，每次只回传当前循环的梯度。
        if b.grad is not None:
            b.data.zero_()
        loss.backward()  # 反向传播
        w.data = w.data - learning_rate * w.grad
        b.data = b.data - learning_rate * b.grad
        print("w,b,loss:\n", w.item(), b.item(), loss.item())

    plt.figure(figsize=(20, 8), dpi=100)
    plt.scatter(x.numpy().reshape(-1), y_true.numpy().reshape(-1))
    y_predict = torch.matmul(x, w) + b
    plt.plot(x.numpy().reshape(-1), y_predict.detach().numpy().reshape(-1), c='r')
    plt.show()


# 使用nn
class Lr(nn.Module):
    def __init__(self):
        super(Lr, self).__init__()  # 必须继承父类
        # torch.nn.Linear(in_features, # 输入的神经元个数
        #            out_features, # 输出神经元个数
        #            bias=True # 是否包含偏置
        #            )
        # 我理解就是输入是1维的，输出也是1维的，因为此时是一个一次的线性的公式
        self.linear = nn.Linear(in_features=1, out_features=1)  # nn.Linear为torch预定义好的线性模型，也被称为全链接层

    def forward(self, x):  # 必须实现forward方法，用于定义神经网络的向前计算过程，即接受输入生产输出的过程。
        out = self.linear(x)
        return out


# 使用pytorch的api
def test2():
    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
    # 定义数据
    x = torch.rand([500, 1]).to(device)
    y_true = x * 3 + 0.8  # 这里我们确定好最终的公式，实际应用里是不知道的，这里是验证最终w是3，b是0.8，不过只能得到一个无限接近的结果

    # 定义模型
    model = Lr().to(device)
    loss_fn = nn.MSELoss(reduction="mean")  # 均方差损失，还有很多其他的损失公式，比如CrossEntropyLoss
    optimizer = optim.SGD(model.parameters(), lr=1e-3)  # 优化器，SGD算法

    # 训练
    for i in range(50000):  # 加大到50000基本就算的很准确了
        y_predict = model(x)  # 调用了__call__方法，call方法又去调用的forward方法，直接就能获取到预测值
        # print(y_predict) # 打印一下看看是不是转换成了GPU的tensor
        loss = loss_fn(y_predict, y_true)  # 计算损失
        optimizer.zero_grad()  # 梯度归零
        loss.backward()  # 梯度计算，反向传播
        optimizer.step()  # 更新梯度
        if i % 100 == 0:
            print("loss:", loss.data)
            result_parameters = list(model.parameters())
            print("w,b:", result_parameters[0].item(), result_parameters[1].item())

    # 模型评估
    model.eval()  # 设置模型为评估模式，即预测模式，保证获取的参数是评估的结果，有时候训练模式和评估模式获取的数据不一样
    predict = model(x)
    predict_data = predict.cpu().data.numpy()
    plt.figure(figsize=(20, 8), dpi=100)
    plt.scatter(x.cpu().data.numpy(), y_true.cpu().data.numpy())
    plt.plot(x.cpu().data.numpy(), predict_data, c='r')
    plt.show()


# https://archive.ics.uci.edu/ml/datasets/SMS+Spam+Collection
# 数据加载
class MyDataset(Dataset):
    def __init__(self, data_path):
        self.lines = open(data_path, encoding='utf-8').readlines()

    def __getitem__(self, index):
        current_line = self.lines[index].strip()
        label = current_line[:4].strip()
        content = current_line[4:].strip()
        return label, content

    def __len__(self):
        return len(self.lines)


def test3():
    my_dataset = MyDataset("SMSSpamCollection")
    print(my_dataset[100])
    print(len(my_dataset))

    # 数据加载器，shuffle加载的时候并发加载，数据都是乱序的
    # batch_size，每一条返回值是多行的合集
    data_loader = DataLoader(dataset=my_dataset, batch_size=5, shuffle=True, num_workers=4)
    for i in data_loader:
        print(i)


# torchvision图形数据集,torchtext文字数据集
def test4():
    # train=True代表下载的是训练集
    train_ds = MNIST(root="./data/", train=True, download=True)
    # 显示图片
    train_ds[0][0].show()
    # 转成tensor
    # 转换形状为(C,H,W)，原始形状为[H,W,C]，等于做一次tensor.permute(2,0,1)
    # (H,W,C)意思为(高，宽，通道数)，黑白图片的通道数只有1，每个像素点的取值为0-255，彩塑图片的通道数为(R,G,B)，三个通道的颜色互相叠加形成颜色
    ret = transforms.ToTensor()(train_ds[0][0])
    print(ret)
    # 图片标准化处理
    # mean均值和std方差的形状要与通道数形状相同
    # 黑白图片的通道数只有1,所以mean和std都只有一个值
    normal = transforms.Normalize(mean=(0.1307,), std=(0.3081,))
    print(normal)


def get_dataloader(train=True, batch_size=128):
    transform_fn = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=(0.1307,), std=(0.3081,))])
    # transform是处理函数
    data_set = MNIST(root="./data/", train=train, download=True, transform=transform_fn)
    data_loader = DataLoader(dataset=data_set, batch_size=batch_size, shuffle=True)
    return data_loader


# 手写一个数字识别
# 最终的预测图片结果是0-9，所以是分类模型
# 我们设计使用三层的神经网络，两个全连接层和一个输出层，第一个全连接层会经过激活函数的处理并交给下一个全连接层，最后输出层输出结果
# 原始的数据形状为[batch_size,1,28,28]，是一个四维的数据，所以要降维变成[batch_size,28,28]，经过第一个连接层输出为[batch_size,28]，这里的28可以是任意值
# 通过激活函数进行处理，因为这是线性变换，所以无法很好的处理数据中的那些非线性特征，于是引入激活函数，提高神经网络的识别能力，比如ReLU激活函数可以帮助神经网络更好地提取数字的轮廓、边缘等特征，从而提高识别的准确率。
# 经过第二个连接层处理，输出的数据形状为[batch_size,10]，因为识别的结果是10个，0-9

class MnistModule(nn.Module):
    def __init__(self):
        super(MnistModule, self).__init__()
        self.fc1 = nn.Linear(1 * 28 * 28, 28)  # 定义第一连接层的线性模型，输入形状是1 * 28 * 28，输出形状是28
        self.fc2 = nn.Linear(28, 10)  # 定义第二连接层的线性模型，输入形状是28，输出形状是10

    def forward(self, x):
        x = x.view(-1, 1 * 28 * 28)  # 先给数据进行变形，-1表示该位置根据后面的形状自行调整，[batch_size,1,28,28]->[batch_size,28,28]
        x = self.fc1(x)  # 第一层神经网络，[batch_size,28,28]->[batch_size,28]
        x = F.relu(x)  # 激活函数不会改变形状，[batch_size,28]->[batch_size,28]
        out = self.fc2(x)  # 第二层神经网络，[batch_size,28]->[batch_size,10]
        return F.log_softmax(out, dim=-1)


def mnist_train(epoch):
    model = MnistModule()
    optimizer = Adam(model.parameters(), lr=0.001)
    # 加载模型
    if os.path.exists("./model/mnist.pkl"):
        model.load_state_dict(torch.load("./model/mnist.pkl"))
    if os.path.exists("./results/optimizer.pkl"):
        optimizer.load_state_dict(torch.load("./model/optimizer.pkl"))

    train_data = get_dataloader(True)

    if not os.path.exists("./model/"):
        os.mkdir("./model/")

    for i in range(epoch):
        for idx, (input, target) in enumerate(train_data):
            optimizer.zero_grad()
            output = model(input)  # 调用模型，获得预测值
            loss = F.nll_loss(output, target)  # nll_loss带权损失，等同于交叉熵损失
            loss.backward()
            optimizer.step()
            if idx % 100 == 0:
                print(loss.item())
            # 保存模型，每训练100个数据就保存一次
            torch.save(model.state_dict(), "./model/mnist.pkl")
            torch.save(optimizer.state_dict(), "./model/optimizer.pkl")


# 模型评估
def mnist_test():
    loss_list = []
    acc_list = []
    model = MnistModule()
    # 加载模型
    model.load_state_dict(torch.load("./model/mnist.pkl"))
    model.eval()  # 设置为评估模式
    test_data = get_dataloader(False, batch_size=128)
    for idx, (input, target) in enumerate(test_data):
        with torch.no_grad():
            output = model(input)
            cur_loss = F.nll_loss(output, target)
            loss_list.append(cur_loss)
            # 计算准确率
            pred = output.max(dim=-1)[-1]
            cur_acc = pred.eq(target).float().mean()  # pred.eq(target)，pred和target是否相等，相等就是预测正确了
            acc_list.append(cur_acc)
            print("平均准确率、平均损失：", np.mean(acc_list), np.mean(loss_list))


if __name__ == '__main__':
    # mnist_train(1)
    mnist_test()
